import pandas as pd
from gofishing.base import Base, TradeType
from gofishing.base import Log
from gofishing.broker import OrderStatus, OrderInfo
from gofishing.broker import Broker
from gofishing.commodity import CommodityBase


class Strategy:

    strategies: dict[str, 'Strategy'] = {}

    def __init__(self, stratID: str, commodity: CommodityBase) -> None:
        if stratID in Strategy.strategies:
            raise Exception('stratID must be unique')
        else:
            self.stratID = stratID
            Strategy.strategies[self.stratID] = self

        self.orderIDs = []
        self.commodity = commodity

        self.long = 0
        self.short = 0
        self.longOpenPrice = 0.0
        self.shortOpenPrice = 0.0

        self.longOpening = 0  # 未 finished 的订单占用
        self.longClosing = 0
        self.shortOpening = 0
        self.shortClosing = 0

        self.record = []

    def _holdingChange(self, orderInfo: OrderInfo) -> None:
        longPre = self.long
        longOpenPricePre = self.longOpenPrice
        shortPre = self.short
        shortOpenPricePre = self.shortOpenPrice
        if orderInfo.volumeLeft:
            # 没有 finished 的 order 不能 _holdingChange
            # 同步部分成交需要考虑重复调用 _holdingChange 的问题
            raise Exception('order not finished')
        if orderInfo.tradeType == TradeType.LONGOPEN:
            self.long += orderInfo.volume
            self.longOpenPrice = \
                (longPre * longOpenPricePre +
                 orderInfo.volume * orderInfo.tradePrice) / self.long
        elif orderInfo.tradeType == TradeType.LONGCLOSE:
            self.long -= orderInfo.volume
            if not self.long:
                self.longOpenPrice = 0.0
        elif orderInfo.tradeType == TradeType.SHORTOPEN:
            self.short += orderInfo.volume
            self.shortOpenPrice = \
                (shortPre * shortOpenPricePre +
                 orderInfo.volume * orderInfo.tradePrice) / self.short
        elif orderInfo.tradeType == TradeType.SHORTCLOSE:
            self.short -= orderInfo.volume
            if not self.short:
                self.shortOpenPrice = 0.0

    def _order(self, tradeType: TradeType,
               limitPrice: float, volume: int) -> None:
        orderID = Broker.order(self.commodity.contractCode,
                               tradeType=tradeType,
                               limitPrice=limitPrice,
                               volume=volume)
        self.orderIDs.append(orderID)

    def _updateOrdering(self) -> None:
        if self.orderIDs:
            orderIDs = self.orderIDs.copy()
            self.longOpening, self.longClosing = 0, 0
            self.shortOpening, self.shortClosing = 0, 0
            # Bug记录：之前用的是“for orderID in self.orderIDs:”，
            # 而“self.orderIDs.remove(orderID)”在遍历的过程中将self.orderIDs改变，
            # 导致漏掉一些orderID。因此，在遍历循环中一定注意遍历对象必须是静止的。
            for orderID in orderIDs:
                orderInfo = Broker.getOrderInfo(orderID)
                if orderInfo.status == OrderStatus.ALIVE:
                    tradeType = orderInfo.tradeType
                    if tradeType == TradeType.LONGOPEN:
                        self.longOpening += orderInfo.volume
                    elif tradeType == TradeType.LONGCLOSE:
                        self.longClosing += orderInfo.volume
                    elif tradeType == TradeType.SHORTOPEN:
                        self.shortOpening += orderInfo.volume
                    elif tradeType == TradeType.SHORTCLOSE:
                        self.shortClosing += orderInfo.volume
                elif orderInfo.status == OrderStatus.FINISHED:
                    self._holdingChange(orderInfo=orderInfo)
                    self._addRecord(orderInfo=orderInfo)
                    self.orderIDs.remove(orderID)
                elif orderInfo.status == OrderStatus.ERROR:
                    self.orderIDs.remove(orderID)

    def _checkBrokerPosition() -> None:
        brokerHolding = Broker.getHolding().copy()
        strategies = Strategy.strategies
        for stratID in strategies:
            strategy = strategies[stratID]
            contractCode = strategy.commodity.contractCode
            if contractCode in brokerHolding:
                brokerHolding[contractCode]['long'] -= strategy.long
                brokerHolding[contractCode]['short'] -= strategy.short
            elif strategy.long or strategy.short:
                raise Exception('Broker-Strategy position not corresponding')
        for contractCode in brokerHolding:
            if brokerHolding[contractCode]['long'] < 0 or \
                    brokerHolding[contractCode]['short'] < 0:
                raise Exception('Broker-Strategy position not corresponding')

    def updateDecision(self) -> None:
        raise NotImplementedError

    def onCloseOfDay(self) -> None:
        raise NotImplementedError

    def update() -> None:
        strategies = Strategy.strategies
        for stratID in strategies:
            strategies[stratID]._updateOrdering()
        Strategy._checkBrokerPosition()
        for stratID in strategies:
            strategies[stratID].updateDecision()
        if Base.isCloseOfDay():
            for stratID in strategies:
                strategies[stratID].onCloseOfDay()

    def open(self, tradeType: TradeType,
             limitPrice: float, volume: int) -> None:
        if tradeType not in [TradeType.LONGOPEN, TradeType.SHORTOPEN]:
            raise Exception('tradeType not statisfied')
        commodity = self.commodity
        margin = volume * \
            (limitPrice * commodity.mult * commodity.marginRate +
             commodity.openCommission)
        commission = volume * commodity.openCommission
        cashNeeded = margin + commission
        available = Broker.getAccountValues().available
        if cashNeeded <= available:
            self._order(tradeType=tradeType,
                        limitPrice=limitPrice,
                        volume=volume)
        else:
            Log.warning('[stratID %s] no enough cash for open' %
                        (self.stratID))

    def close(self, tradeType: TradeType,
              limitPrice: float, volume: int) -> None:
        if tradeType not in [TradeType.LONGCLOSE, TradeType.SHORTCLOSE]:
            raise Exception('tradeType not statisfied')

        longFree = self.long - self.longClosing
        shortFree = self.short - self.shortClosing
        if tradeType == TradeType.LONGCLOSE:
            if longFree >= volume:
                self._order(tradeType=tradeType,
                            limitPrice=limitPrice,
                            volume=volume)
            else:
                Log.warning(
                    '[stratID %s] no enough long position for close' %
                    (self.stratID))
        elif tradeType == TradeType.SHORTCLOSE:
            if shortFree >= volume:
                self._order(tradeType=tradeType,
                            limitPrice=limitPrice,
                            volume=volume)
            else:
                Log.warning(
                    '[stratID %s] no enough short position for close' %
                    (self.stratID))

    def isNoPosition(self) -> bool:
        if not self.orderIDs and \
                not self.long and not self.short and \
                not self.longOpening and not self.longClosing and \
                not self.shortOpening and not self.shortClosing:
            return True
        else:
            return False

    def _addRecord(self, orderInfo: OrderInfo) -> None:
        rowData = [orderInfo.tradeTime,
                   orderInfo.insertTime,
                   self.stratID,
                   self.commodity.commodityID,
                   self.commodity.contractCode,
                   orderInfo.orderID,
                   orderInfo.tradeType,
                   orderInfo.limitPrice,
                   orderInfo.tradePrice,
                   orderInfo.volume - orderInfo.volumeLeft,
                   self.long, self.short,
                   self.longOpenPrice, self.shortOpenPrice,
                   0.0, orderInfo.commission]
        self.record.append(rowData)

    def composeRecord(self) -> pd.DataFrame:
        columns = ['tradeTime', 'insertTime',
                   'stratID', 'commodityID', 'contractCode',
                   'orderID', 'tradeType', 'limitPrice', 'tradePrice',
                   'volume',
                   'long', 'short', 'longOpenPrice', 'shortOpenPrice',
                   'rawProfit', 'commission']
        # 计算 rawProfit 列
        mult = self.commodity.mult
        record = pd.DataFrame(data=self.record, columns=columns)
        long, short = record['long'], record['short']
        longOpenPrice, shortOpenPrice = \
            record['longOpenPrice'], record['shortOpenPrice']
        tradePrice = record['tradePrice']
        indicePre = 0
        for indice in record.index:
            if indice != 0:
                longDiff = long[indice] - long[indicePre]
                if longDiff < 0:
                    record.at[indice, 'rawProfit'] += \
                        (mult * (-longDiff) *
                         (tradePrice[indice] - longOpenPrice[indicePre]))
                shortDiff = short[indice] - short[indicePre]
                if shortDiff < 0:
                    record.at[indice, 'rawProfit'] += \
                        (mult * (-shortDiff) *
                         (shortOpenPrice[indicePre] - tradePrice[indice]))
                indicePre = indice
        return record

    def filterRecord(self, tradeType: TradeType) -> pd.DataFrame:
        rec = self.composeRecord()
        return rec[rec['tradeType'] == tradeType].reset_index(drop=True)
